nemo_rl/models/generation/vllm/vllm_backend.py (1)

234-234: Replace lambda with local function per style guidelines.

Static analysis flagged E731: assigning a lambda expression. Use a def statement instead.

♻️ Suggested fix

-        load_model_weight_func = lambda x: self._load_model_weights(x, self.model_runner)
+        def load_model_weight_func(x):
+            self._load_model_weights(x, self.model_runner)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@nemo_rl/models/generation/vllm/vllm_backend.py` at line 234, Replace the
lambda assignment to load_model_weight_func with a local def function to satisfy
style rules: create a small local function (e.g., def
load_model_weight_func(path): return self._load_model_weights(path,
self.model_runner)) and use that function in place of the lambda; keep the same
name load_model_weight_func and ensure it calls self._load_model_weights with
the same arguments (path and self.model_runner) so behavior does not change.

nemo_rl/models/generation/vllm/quantization/fp8_train_utils.py (3)

57-62: Consider zero-padding instead of edge-value padding.

Using data_hp[-1, -1] for padding could propagate anomalous values (NaN, Inf, or outliers) if they happen to be at the tensor edge. Zero-padding is more conventional and predictable for quantization.

🔧 Suggested fix

         data_hp = torch.nn.functional.pad(
-            data_hp, (0, pad1, 0, pad0), mode="constant", value=data_hp[-1, -1]
+            data_hp, (0, pad1, 0, pad0), mode="constant", value=0.0
         )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@nemo_rl/models/generation/vllm/quantization/fp8_train_utils.py` around lines
57 - 62, Replace edge-value padding with zero-padding in the block-alignment
logic: when computing pad0/pad1 and calling torch.nn.functional.pad on data_hp
(the block-padding branch that checks data_hp.shape[...] % block_size...),
change the pad fill from value=data_hp[-1, -1] to value=0 so that padding uses
zeros instead of the tensor's edge element; keep the same pad tuple order and
mode="constant" so only the fill value changes.

---

69-69: Undocumented square-block constraint.

The function signature accepts weight_block_size: list[int] suggesting arbitrary block dimensions, but this assertion enforces square blocks. Consider documenting this restriction in the docstring or simplifying the signature if non-square blocks are never intended.

📝 Document the restriction in docstring

     Args:
         data_hp: 2-D high-precision weight tensor (any float dtype).
-        weight_block_size: [block_rows, block_cols], e.g. [128, 128].
+        weight_block_size: [block_rows, block_cols], e.g. [128, 128].
+            Note: block_rows must equal block_cols (square blocks only).
         use_pow2_scale: If True, round scale factors to powers of two.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@nemo_rl/models/generation/vllm/quantization/fp8_train_utils.py` at line 69,
The code currently asserts block_size0 == block_size1 (from the assert
block_size0 == block_size1) while the function accepts weight_block_size:
list[int], which implies non-square blocks are allowed; update the function’s
docstring (the function that takes weight_block_size) to state explicitly that
only square blocks are supported (i.e., weight_block_size must be [N, N]) or, if
non-square blocks are never intended, simplify the signature to accept a single
int for block_size and remove the list ambiguity; reference the
weight_block_size parameter and the assert block_size0 == block_size1 when
documenting the constraint.

---

52-52: Consider using explicit exception instead of assert for input validation.

Assertions can be disabled with Python's -O flag, which would bypass this validation in optimized runs. For production code, explicit exceptions are more reliable.

🔧 Suggested fix

-    assert len(data_hp.shape) == 2, "Only 2-D input tensor is supported"
+    if len(data_hp.shape) != 2:
+        raise ValueError("Only 2-D input tensor is supported")

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@nemo_rl/models/generation/vllm/quantization/fp8_train_utils.py` at line 52,
Replace the runtime-only assert on data_hp with an explicit input validation
that always runs: check if len(data_hp.shape) != 2 and raise a ValueError with
the message "Only 2-D input tensor is supported"; update the check around the
data_hp usage in fp8_train_utils.py (the line currently doing `assert
len(data_hp.shape) == 2, ...`) to this explicit conditional to ensure validation
cannot be skipped under -O.